using System.Net.WebSockets;
using System.Collections.Concurrent;
using System.Text;
using IMServer.WebSocketServices;

public class WebSocketManager
{
    private static readonly Lazy<WebSocketManager> _instance = 
        new Lazy<WebSocketManager>(() => new WebSocketManager());
    
    private readonly ConcurrentDictionary<string, WebSocketConnection> _connections;
    private readonly CancellationTokenSource _shutdownTokenSource;

    public static WebSocketManager Instance => _instance.Value;
    
    private WebSocketManager()
    {
        _connections = new ConcurrentDictionary<string, WebSocketConnection>();
        _shutdownTokenSource = new CancellationTokenSource();
    }

    // 添加关闭方法
    public async Task ShutdownAsync()
    {
        try
        {
            _shutdownTokenSource.Cancel();
            
            // 关闭所有连接
            var closeTasks = _connections.Values.Select(async connection =>
            {
                try
                {
                    if (connection.WebSocket.State == WebSocketState.Open)
                    {
                        await connection.WebSocket.CloseAsync(
                            WebSocketCloseStatus.NormalClosure,
                            "Server is shutting down",
                            CancellationToken.None);
                    }
                }
                catch (Exception ex)
                {
                    Console.WriteLine($"Error closing connection: {ex.Message}");
                }
            });

            // 等待所有连接关闭，设置超时时间
            await Task.WhenAll(closeTasks).WaitAsync(TimeSpan.FromSeconds(5));
            
            // 清空连接字典
            _connections.Clear();
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Error during shutdown: {ex.Message}");
        }
    }

    // 修改 HandleConnectionAsync 方法，添加取消令牌
    public async Task HandleConnectionAsync(HttpContext context, string userId)
    {
        if (!context.WebSockets.IsWebSocketRequest)
        {
            context.Response.StatusCode = StatusCodes.Status400BadRequest;
            return;
        }

        using var webSocket = await context.WebSockets.AcceptWebSocketAsync();
        var connection = new WebSocketConnection(userId, webSocket);

        try
        {
            _connections.TryAdd(connection.ConnectionId, connection);
            await KeepConnectionAliveAsync(connection);
        }
        catch (OperationCanceledException)
        {
            // 正常关闭，不需要处理
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Connection error: {ex.Message}");
        }
        finally
        {
            _connections.TryRemove(connection.ConnectionId, out _);
        }
    }

    // 修改 KeepConnectionAliveAsync 方法，使用取消令牌
    private async Task KeepConnectionAliveAsync(WebSocketConnection connection)
    {
        var buffer = new byte[1024 * 4];
        try
        {
            while (!_shutdownTokenSource.Token.IsCancellationRequested &&
                   connection.WebSocket.State == WebSocketState.Open)
            {
                var result = await connection.WebSocket.ReceiveAsync(
                    new ArraySegment<byte>(buffer),
                    _shutdownTokenSource.Token);

                if (result.MessageType == WebSocketMessageType.Close)
                    break;

                if (result.Count > 0)
                {
                    var message = Encoding.UTF8.GetString(buffer, 0, result.Count);
                    await WebSocketMessageHandle.HandleMessage(message);
                }
            }
        }
        finally
        {
            if (connection.WebSocket.State == WebSocketState.Open)
            {
                await connection.WebSocket.CloseAsync(
                    WebSocketCloseStatus.NormalClosure,
                    "Closing connection",
                    CancellationToken.None);
            }
        }
    }

    // 向指定用户发送消息
    public async Task<bool> SendMessageToUserAsync(string userId, string message)
    {
        var connection = _connections.Values.FirstOrDefault(c => c.UserId == userId);
        if (connection?.WebSocket.State == WebSocketState.Open)
        {
            try 
            {
                var bytes = Encoding.UTF8.GetBytes(message);
                await connection.WebSocket.SendAsync(
                    new ArraySegment<byte>(bytes),
                    WebSocketMessageType.Text,
                    true,  // endOfMessage: true 表示这是完整的消息
                    CancellationToken.None);
                return true;
            }
            catch
            {
                return false;
            }
        }
        return false;
    }
    
    // 检查某个用户是否在线
    public bool IsUserOnline(string userId)
    {
        return _connections.Values
            .Any(c => c.UserId == userId && c.WebSocket.State == WebSocketState.Open);
    }
}